/*
 * Public Credit Information Sharing Platform(PCISP).
 * Copyright (c) 2019-2020 jsxzkingdee and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 * Please visit https://www.zhengxinsoft.com if you need additional information or have any questions.
 */

package com.platform.core.limit;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import com.platform.comm.exceptions.RestApiException;
import com.platform.comm.util.HttpUtil;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;

import javax.annotation.Nullable;
import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

/**
 * Description: 自定义限流切面
 *
 * @author libin created on 2020/5/8.
 */
@SuppressWarnings("UnstableApiUsage")
@Component
@Aspect
public class LimitAspect {

    /**
     * 根据IP分不同的令牌桶, 每天自动清理缓存
     */
    private static LoadingCache<String, RateLimiter> caches = CacheBuilder.newBuilder()
            .maximumSize(1000)
            .expireAfterWrite(1, TimeUnit.DAYS)
            .build(new CacheLoader<String, RateLimiter>() {
                @Override
                public RateLimiter load(@Nullable String key) {
                    // 新的IP初始化 每秒只发出5个令牌
                    return RateLimiter.create(5);
                }
            });

    /**
     * 切点  限流
     */
    @Pointcut("@annotation(com.platform.core.limit.CustomLimit)")
    public void customAspect() {

    }

    @Around("customAspect()")
    public Object around(ProceedingJoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        CustomLimit limitAnnotation = method.getAnnotation(CustomLimit.class);
        CustomLimit.LimitType limitType = limitAnnotation.limitType();
        String key = limitAnnotation.key();
        Object obj;
        try {
            if (limitType.equals(CustomLimit.LimitType.IP)) {
                key = HttpUtil.getIpAddress();
            }
            RateLimiter rateLimiter = caches.get(key);
            boolean flag = rateLimiter.tryAcquire();
            if (flag) {
                obj = joinPoint.proceed();
            } else {
                throw new RestApiException("您访问的太频繁了");
            }
        } catch (Throwable e) {
            throw new RestApiException("您访问的太频繁了");
        }
        return obj;
    }
}
